!
!  Robust mean
!
!  Copyright © 2011 F.Hroch (hroch@physics.muni.cz)
!
!  This file is part of Munipack.
!
!  Munipack is free software: you can redistribute it and/or modify
!  it under the terms of the GNU General Public License as published by
!  the Free Software Foundation, either version 3 of the License, or
!  (at your option) any later version.
!  
!  Munipack is distributed in the hope that it will be useful,
!  but WITHOUT ANY WARRANTY; without even the implied warranty of
!  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
!  GNU General Public License for more details.
!  
!  You should have received a copy of the GNU General Public License
!  along with Munipack.  If not, see <http://www.gnu.org/licenses/>.
!

module robustmean

  implicit none

  ! robust estimator of arithmetical mean by
  ! Hogg in Launer, Wilkinson: Robustness in Statistics

  ! print debug informations?
  logical, parameter, private :: print = .false.

  ! precision of real numbers
  integer, parameter, private :: rp = selected_real_kind(15)
  
  interface rmean
     module procedure rmean_double, rmean_single
  end interface rmean

contains

  ! rmean is a basic user function

  subroutine rmean_double(x,t,dt)

    ! simple interface for rmean1

    use rfun
    
    real(rp), dimension(:),intent(in) :: x(:)
    real(rp), intent(out) :: t,dt
    integer :: istat

    call rmean1(x,t,dt,tukey,dtukey,1e-15_rp,10,istat)
    
  end subroutine rmean_double

  subroutine rmean_single(x,t,dt)

    ! this is little bit slower, but the most simple solution

    use rfun

    real, dimension(:),intent(in) :: x(:)
    real, intent(out) :: t,dt
    real(rp), dimension(:), allocatable :: y(:)
    real(rp) :: d,dd
    integer :: istat

    allocate(y(size(x)))
    y = x
    call rmean1(y,d,dd,tukey,dtukey,1e-15_rp,10,istat)
    deallocate(y)
    t = d
    dt = dd

  end subroutine rmean_single

  !-------------------------------------------------------------------

  subroutine medmad(x,med,mad)

    ! estimate median and median of absoluted deviations

    real(rp), dimension(:),intent(in) :: x(:)
    real(rp), intent(out) :: med,mad
    integer :: n
    real(rp), dimension(:),allocatable :: y

    n = size(x)
    call qmed(x,n/2,med)
    ! this is right just only for odd sequences, but
    ! we're ignore it because one is just only an initial  estimate

    allocate(y(n))
    y = abs(x - med)
    call qmed(y,n/2+1,mad)
    deallocate(y)

    if( print ) write(*,*) "med, mad : ",med, mad

  end subroutine medmad

  !-----------------------------------------------------------------------

  subroutine rmean1(x,t,dt,psi,dpsi,rel,maxit,istat)

    ! solves equation for minimization with robust
    ! function by the Newton's iteraction method

    interface
       function psi(x)
         integer, parameter :: rp = selected_real_kind(15)
         real(rp) :: psi
         real(rp),intent(in) :: x
       end function psi
       function dpsi(x)
         integer, parameter :: rp = selected_real_kind(15)
         real(rp) :: dpsi
         real(rp),intent(in) :: x
       end function dpsi
    end interface

    real(rp), dimension(:), intent(in) :: x
    real(rp), intent(out) :: t,dt
    real(rp), intent(in) :: rel
    integer, intent(in) :: maxit
    integer, intent(out) :: istat

    integer :: n,i,it
    real(rp) :: d,r,s,s2,sum1,sum2,sum3,p

    istat = 0; dt = -1.0_rp
    n = size(x)

    if( n < 1 ) then                     ! a few data
       istat = 1
       t = 0.0_rp
    endif

    if( n == 1 ) then               ! only one point, but correct case
       t = x(1)
       dt = 0.0_rp
       return
    endif

    ! initial values    
    call medmad(x,t,s)
    s = s/0.6745_rp
    s2 = s**2

    if( abs(s) < epsilon(s) )then    ! identical points on input
       t = x(1)
!       write(*,*) x(1:n),t,s
       dt = 0.0_rp
       return
    endif

    do it = 1,maxit

       ! corrector's estimation 
       sum1 = 0.0_rp; sum2 = 0.0_rp; sum3= 0.0_rp;
       do i = 1,n
          r = (x(i) - t)/s
          p = psi(r)
          sum1 = sum1 + p
          sum2 = sum2 + dpsi(r)
          sum3 = sum3 + p**2
       enddo
       if( abs(sum2) < epsilon(sum2) )then
          istat = 2               ! ?
          return
       endif
!       write(*,*) sum1,sum2,sum3
       d = s*sum1/sum2
       t = t + d
       dt = s2*n/(n-1)*sum3/sum2**2

       ! exit of iterations: 
       if( it > 2 .and. (d**2 < rel*dt .or. abs(d) < 10.0*epsilon(d)) ) exit

       if( print ) write(*,*) "mean, increment, sigma: ",t,d,sqrt(dt)
    enddo

    ! estimation pf deviation
    sum1 = 0.0_rp; sum3 = 0.0_rp
    do i = 1, n
       r = (x(i) - t)/s
       sum1 = sum1 +  psi(r)**2
       sum3 = sum3 + dpsi(r)
    enddo
    if( sum3 == 0.0_rp )then
       istat = 3                ! ?
       return
    endif
    dt = s2*n/(n-1)*sum1/sum3**2
    dt = sqrt(dt)

  end subroutine rmean1



  !-------------------------------------

  ! Quick median by Wirth: Algorith + data structures = programs

  subroutine qmed(b,k,x)

    integer, intent(in) :: k
    real(rp),dimension(:),intent(in) :: b
    real(rp),dimension(:),allocatable :: a
    real(rp) :: w,x
    integer :: n,l,r,i,j

    n = size(b)
    allocate(a(n))
    a = b

    l = 1
    r = n
    do while( l < r )
       x = a(k)
       i = l
       j = r
       do
          do while( a(i) < x )
             i = i + 1
          enddo
          do while( x < a(j) )
             j = j - 1
          enddo
          if( i <= j ) then
             w = a(i)
             a(i) = a(j)
             a(j) = w
             i = i + 1
             j = j - 1
          endif
          if( i > j ) exit
       enddo
       if( j < k ) l = i
       if( k < i ) r = j
    enddo
    deallocate(a)

  end subroutine qmed


end module robustmean
